package com.mazaiting.cryption.sm.sm4;

import org.bouncycastle.util.encoders.Base64;

import java.nio.charset.StandardCharsets;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * SM4 工具类
 * 源码: https://github.com/SmartXiaoMing001/Chinese-Cipher-Of-SM2-SM3-SM4
 */
public class SM4Util {

    private SM4Util() {
    }

    /**
     * 加密
     *
     * @param key  密钥
     * @param data 待加密数据
     * @return Base64编码加密数据
     */
    public static String encrypt(String key, String data) {
        return ecbEncryptStr(Base64.decode(key), data.getBytes());
    }

    /**
     * 解密
     *
     * @param key  密钥
     * @param data 带解密数据
     * @return 明文
     */
    public static String decrypt(String key, String data) {
        return ecbDecryptStr(key, data);
    }


    /**
     * ECB 加密数据
     *
     * @param keyBytes 密钥
     * @param data     待加密数据数据
     * @return Base64 编码后的数据
     */
    public static String ecbEncryptStr(byte[] keyBytes, byte[] data) {
        String cipherText = "";
        try {
            // 加密
            byte[] encrypted = ecbEncrypt(keyBytes, data);
            if (null != encrypted) {
                // BASE64
                cipherText = new String(Base64.encode(encrypted), StandardCharsets.UTF_8);
                if (cipherText.trim().length() > 0) {
                    Pattern p = Pattern.compile("\\s*|\t|\r|\n");
                    Matcher m = p.matcher(cipherText);
                    cipherText = m.replaceAll("");
                }
            }
        } catch (Exception ignored) {
        }
        return cipherText;
    }

    /**
     * ECB 加密
     *
     * @param keyBytes 密钥字节数组
     * @param data     字节数组
     * @return 加密后的字节数组
     */
    public static byte[] ecbEncrypt(byte[] keyBytes, byte[] data) {
        try {
            SM4Context ctx = new SM4Context(SM4.SM4_ENCRYPT);
            SM4 sm4 = new SM4();
            sm4.sm4SetKeyEnc(ctx, keyBytes);
            // 加密
            return sm4.sm4CryptEcb(ctx, data);
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * ECB 解密字符串
     *
     * @param key  密钥
     * @param data 待解密数据
     * @return 明文
     */
    public static String ecbDecryptStr(String key, String data) {
        try {
            byte[] keyBytes = Base64.decode(key);
            byte[] dataBytes = Base64.decode(data);
            byte[] resBytes = ecbDecrypt(keyBytes, dataBytes);
            if (null != resBytes) {
                return new String(resBytes, StandardCharsets.UTF_8);
            }
        } catch (Exception ignored) {
        }
        return null;
    }

    /**
     * ECB 解密
     *
     * @param keyBytes  密钥字节数组
     * @param dataBytes 待解密数据字节数组
     * @return 明文字节数组
     */
    public static byte[] ecbDecrypt(byte[] keyBytes, byte[] dataBytes) {
        try {
            SM4Context ctx = new SM4Context(SM4.SM4_DECRYPT);
            SM4 sm4 = new SM4();
            sm4.sm4SetKeyDec(ctx, keyBytes);
            return sm4.sm4CryptEcb(ctx, dataBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * CBC 加密
     *
     * @param keyBytes 待加密数据
     * @param ivBytes  混淆字节数组
     * @param data     待加密字节数组
     * @return Base64 编码后的数据
     */
    public static String cbcEncryptStr(byte[] keyBytes, byte[] ivBytes, byte[] data) {
        String cipherText = "";
        try {
            // 加密
            byte[] encrypted = cbcEncrypt(keyBytes, ivBytes, data);
            if (null != encrypted) {
                // BASE64
                cipherText = new String(Base64.encode(encrypted), StandardCharsets.UTF_8);
                if (cipherText.trim().length() > 0) {
                    Pattern p = Pattern.compile("\\s*|\t|\r|\n");
                    Matcher m = p.matcher(cipherText);
                    cipherText = m.replaceAll("");
                }
            }
        } catch (Exception ignored) {
        }
        return cipherText;
    }

    /**
     * CBC 加密
     *
     * @param keyBytes  待加密数据
     * @param ivBytes   混淆字节数组
     * @param dataBytes 待加密数组
     * @return 加密后数组
     */
    private static byte[] cbcEncrypt(byte[] keyBytes, byte[] ivBytes, byte[] dataBytes) {
        try {
            SM4Context ctx = new SM4Context(SM4.SM4_ENCRYPT);
            SM4 sm4 = new SM4();
            sm4.sm4SetKeyEnc(ctx, keyBytes);
            return sm4.sm4CryptCbc(ctx, ivBytes, dataBytes);
        } catch (Exception ignored) {
        }
        return null;
    }

    /**
     * CBC 解密
     *
     * @param key  密钥
     * @param iv   IV填充
     * @param data 待解密数据
     * @return 明文
     */
    public static String cbcDecryptStr(String key, String iv, String data) {
        try {
            byte[] keyBytes = Base64.decode(key);
            byte[] ivBytes = Base64.decode(iv);
            byte[] dataBytes = Base64.decode(data);
            byte[] resBytes = cbcDecrypt(keyBytes, ivBytes, dataBytes);
            if (null != resBytes) {
                return new String(resBytes, StandardCharsets.UTF_8);
            }
        } catch (Exception ignored) {
        }
        return null;
    }

    /**
     * CBC 解密
     *
     * @param keyBytes  密钥数组
     * @param ivBytes   IV 数组
     * @param dataBytes 待解密数组
     * @return 明文数组
     */
    private static byte[] cbcDecrypt(byte[] keyBytes, byte[] ivBytes, byte[] dataBytes) {
        try {
            SM4Context ctx = new SM4Context(SM4.SM4_DECRYPT);
            SM4 sm4 = new SM4();
            sm4.sm4SetKeyDec(ctx, keyBytes);
            return sm4.sm4CryptCbc(ctx, ivBytes, dataBytes);
        } catch (Exception ignored) {
        }
        return null;
    }
}